# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections
import json
import logging

from absl import app, flags

from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.experimental.xguo.fruit.grape.price_graph import PriceGraph, OrderSide
from coin.experimental.xguo.fruit.base.event_clock import (
    event_clock_scheduler,
    event_clock,
)
from coin.experimental.xguo.fruit.base.strategy_sim_base import (
    BaseSimStrategy,
    BaseStrategyConfig,
)
from coin.strategy.mm.simple_sim.executors import PassiveSimExecutor
from coin.experimental.xguo.fruit.util.util import FillStats

FLAGS = flags.FLAGS


class TriMMSimStrategy(BaseSimStrategy):
  def __init__(self, config):
    super().__init__(config)
    self._pricer = PriceGraph(self.product_list, ignore_fee=True)
    event_clock_scheduler.add_periodic_task(1800, self.log_sim_time)
    event_clock_scheduler.add_periodic_task(300, self.pnl_snapshot)
    event_clock_scheduler.add_periodic_task(3600, self.save_sim_result)
    self.pnl_report = []
    self._fill_stats = collections.defaultdict(lambda: None)

  def on_book(self, product, book):
    timestamp = book.timestamp
    self._pricer.update(book)
    if not self._pricer.is_ready():
      return

    if self._fill_stats[product] is None:
      self._fill_stats[product] = FillStats(product, fee_bp=1)
      event_clock_scheduler.add_periodic_task(
          period=30,
          job_func=self._fill_stats[product].print,
          msg='Sim PNL from fill',
      )
    self._fill_stats[product].update_book(book)

    fee_bp = get_fee_rate(product)
    maker_fee = fee_bp[0] * 1e-4
    taker_fee = fee_bp[1] * 1e-4
    edge_bp = fee_bp[0] * 0 + 2
    if product not in self.executors:
      config = self.strategy_config.executor_config[product]['proto_config']
      min_pos = config.min_pos
      max_pos = config.max_pos
      lot_size = config.lot_size
      min_fire_period = config.order_update_period_sec * 1e9
      executor = PassiveSimExecutor(
          lot_size=lot_size,
          min_pos=min_pos,
          max_pos=max_pos,
          maker_fee=maker_fee,
          taker_fee=taker_fee,
          execution_delay=100,
          post_only=False,
          min_fire_period=min_fire_period,
      )

      # noinspection PyUnusedLocal
      def sim_fill_callback(price, qty, side, maker, ts):
        self._fill_stats[product].update_sim(price, qty, side)

      executor.fill_callback = sim_fill_callback
      self.executors[product] = executor

    self.executors[product].update_book(book)
    orders = self._pricer.gen_mm_orders(edge_bp=edge_bp)

    for prod, order in orders.items():
      if prod == product:
        continue

      executor = self.executors.get(prod)
      if executor:
        executor.post(
            timestamp=timestamp,
            buy_price=order[OrderSide.BUY].price,
            sell_price=order[OrderSide.SELL].price,
        )
    self._logger.debug('Submit order: %s', event_clock.now())

  def on_trade(self, product, trade_event):
    executor = self.executors.get(product)
    if executor:
      executor.update_trade(trade_event)

  def pnl_snapshot(self):
    timestamp = event_clock.timestamp
    result = {
        'timestamp': timestamp,
        'human_readable_timestamp': str(event_clock.now()),
        'data': {},
    }
    for product, executor in self.executors.items():
      result['data'][product.symbol] = executor.get_summary()
    print(json.dumps(result, indent=2))
    self.pnl_report.append(result)

  def save_sim_result(self):
    filename = self.strategy_config.sim_config.get('result_file')
    if not filename:
      filename = str(event_clock.timestamp) + '.json'
    with open(filename, 'w') as out_file:
      json.dump(self.pnl_report, out_file)


def main(_):
  config = BaseStrategyConfig.from_file(FLAGS.config)
  strategy = TriMMSimStrategy(config)
  strategy.run()
  strategy.save_sim_result()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
